import numpy as np
import pandas as pd
import re
import os
import json
import warnings
import colorama

warnings.filterwarnings("ignore")

from get_number_in_filename import Get_number_in_filename

# TEXT_COLOR_RED = '\033[31m'
# TEXT_COLOR_GREEN = '\033[32m'
# TEXT_COLOR_YELLOW = '\033[33m'
# TEXT_COLOR_END = '\033[0m'

colorama.init(autoreset=True)
TEXT_COLOR_RED = colorama.Fore.RED
TEXT_COLOR_GREEN = colorama.Fore.GREEN
TEXT_COLOR_YELLOW = colorama.Fore.YELLOW
TEXT_COLOR_END = colorama.Fore.RESET
root_path=os.path.join(os.getcwd(),".\\Data_process\\Parse_Setfos")#为了支持auto-process.py的调用

def coef_json2dict(coef_file_name):
    '''
    input: json_file: {"CE": {"0": 1, "15": 1, "30": 1, "45": 1, "60": 1}}
    return: dict: {'CE_0': 1, 'CE_15': 1, 'CE_30': 1, 'CE_45': 1, 'CE_60': 1}
    '''
    def rawdict2dict(rawdict):
        coef_dict = {}
        for sweep in rawdict:
            for json_angle in rawdict[sweep]:
                coef_dict[f'{sweep}_{json_angle}'] = rawdict[sweep][json_angle]
        return coef_dict

    try:
        exist_json_dict = json.load(open(os.path.join(root_path,coef_file_name), 'r'))
        return rawdict2dict(exist_json_dict)
    except:
        print('Can not find ' + TEXT_COLOR_RED + coef_file_name + TEXT_COLOR_END)
        sample_coef_dict = {"CE": {"0": 1, "15": 1, "30": 1, "45": 1, "60": 1}}
        with open(os.path.join(root_path, coef_file_name), 'w') as f:
            json.dump(sample_coef_dict, f)
        print('Generated a ' + TEXT_COLOR_GREEN + coef_file_name + TEXT_COLOR_END)
        return rawdict2dict(sample_coef_dict)

def txt2df(lines):
    global angles
    ncond = 0
    mark_line = 0
    for i, line in enumerate(lines):
        matched = re.match('# Sweep ([1-9][0-9]*):', line)
        if matched:
            ncond = int(matched.groups()[0])
        else:
            if ncond > 0:
                mark_line = i
                break

    if ncond < 1:
        raise IOError('No Sweep Conditions')

    columns = [x.strip() for x in re.findall('([^#()]+)\(.*?\)', lines[mark_line + 1])]#['HTL.d', 'Cathode.d', 'LiF.d', 'SiON-B.d', 'CIE_x,R', 'CIE_y,R', 'CIE_Y,R', 'CIE_x,T', 'CIE_y,T', 'CIE_Y,T', 'CRI transmitted', 'SEIR', 'SIA', 'Rref', 'Tref', 'Rad', 'Rad.s', 'Rad.p', 'Luminance', 'Current Efficiency', 'Luminous Efficacy', 'CIE_x', 'CIE_y', 'CIE_Y', 'CIE_x,0', 'CIE_y,0', 'CRI', 'CCT', 'Delta_uv', 'CIE_L', 'CIE_a', 'CIE_b', 'CIE_u', 'CIE_v', 'Blue index', 'Peak intensity', 'Peak wavelength', 'Peak intensity FWHM', 'Peak intensity.s', 'Peak wavelength.s', 'Peak intensity.p', 'Peak wavelength.p']
    # print(columns)
    
    columns[columns.index('Current Efficiency')] = 'CE'
    columns[columns.index('CIE_x')] = 'CIEx'
    columns[columns.index('CIE_y')] = 'CIEy'
    columns[columns.index('Peak wavelength')] = 'Peak'

    raw_datas = np.array([line.strip().split() for line in lines[mark_line + 2:] if len(line) > 1], dtype=np.double)
    # print(raw_datas)
    raw_df = pd.DataFrame(raw_datas, columns=columns)
    # raw_df.insert(0, 'AdditionalInfo', additional_string)

    # print(raw_df)
    # raw_df['Filename'] = 'filename'


    if columns[ncond - 1] == 'EmissionAngle':
        angles = raw_df.loc[:, 'EmissionAngle'].unique().astype(np.int64)
        new_columns = columns[:ncond - 1]
        angle_datas = []
        max_rows = len(raw_df) // len(angles)
        for i, angle in enumerate(angles):
            angle_part = raw_df[(raw_df.loc[:, 'EmissionAngle'] - angle).abs() < 1E-5]
            if i == 0:
                angle_datas.append(angle_part.loc[:, columns[:ncond - 1]].values[:max_rows, :])
            new_columns.extend(['%s_%d' % (x, angle,) for x in ['CE', 'CIEx', 'CIEy', 'Peak']])
            angle_datas.append(angle_part.loc[:, ['CE', 'CIEx', 'CIEy', 'Peak']].values[:max_rows, :])
        angle_datas = np.concatenate(angle_datas,axis=1)
        # out_df = pd.DataFrame(angle_datas, columns=new_columns).assign(AdditionalInfo=additional_string)
        out_df = pd.DataFrame(angle_datas, columns=new_columns)
    else:
        # out_df = raw_df.loc[:, columns[:ncond] + ['CE', 'CIEx', 'CIEy', 'Peak']].assign(AdditionalInfo=additional_string)
        out_df = raw_df.loc[:, columns[:ncond] + ['CE', 'CIEx', 'CIEy', 'Peak']]
        # print(out_df)
        # out_df = raw_df
    # print("***")
    # print(columns[:ncond])
    # print("***")
    return out_df,columns[:ncond]






def fout(input_file_name,process_name,out_format):
    info_list=input_file_name.split(".")[0].split("-")[0].split("_")
    fout_prefix = f'{process_name}_'+"_".join(info_list[1:])
    ntry = 0
    while True:
        if ntry == 0:
            fout = fout_prefix + f'.{out_format}'
        else:
            fout = fout_prefix + '-%d.%s' % (ntry,out_format)
        if os.access(os.path.join(root_path,fout), os.F_OK):
            ntry += 1
        else:
            break
    return fout

def main(lines_list,coef_config,string_list):
    pd_list = []
    angles = []
    string1_len=0
    string_name=Get_number_in_filename(string_list)
    print("***********")
    print(string_name)
    print("***********")
    out_df ,out_cond= txt2df(lines_list[0],)
    print('填入第一个合并填充内容')
    # string1=int(string_name[0])
    string1=string_name[0]
    print(string1)
    for i in range(0,len(string1)):
        print(f'Luminescent_layer(NK{i})')
        out_df[f'Luminescent_layer(NK{i})'] = int(string1[i])
    
    if len(lines_list) > 1:
        j=1
        for i in lines_list[1:]:
            df,cond = txt2df(i,)

            string3=f'填入第{j+1}个合并填充内容'
            print(string3)
            # string2=input()
            if j>len(string_name):
                break
            # string2=int(string_name[j])
            # print(string2)
            # df['Luminescent_layer(NK)'] = string2
            string2=string_name[j]
            print(string2)
            for k in range(0,len(string2)):
                print(f'Luminescent_layer(NK{k})')
                df[f'Luminescent_layer(NK{k})'] = int(string2[k])
            cols_contige_bool=(df.columns.tolist()==out_df.columns.tolist())
            print(cols_contige_bool)
            if cols_contige_bool:
                out_df=pd.concat([out_df,df])
                j=j+1
            else:
                raise ValueError ("特征不一致，无法合并")
    # out_df.columns = [f"Setfos_{col}" for col in out_df.columns]
    # coef系数

    # out_df = out_df[['Luminescent_layer(NK0)'] +['Luminescent_layer(NK1)'] +list(out_df.columns[:-1])]
    row_names = [f'Luminescent_layer(NK{i})' for i in range(0,len(string1))]
    # print(f'row_names={row_names}')
    string1_len=len(string1)
    new_column_order = row_names + list(out_df.columns[:-string1_len])
    # print(f'new_column_order={new_column_order}')
    out_df = out_df[new_column_order]
    # print(out_df.columns)

    # print(row_name)
    # out_df=out_df[row_name]




    coef_dict = coef_json2dict(coef_config)
    coef_cols = coef_dict.keys()

    matched_coef_cols = list(set(out_df.columns) & set(coef_cols))
    matched_coef_dict = dict([(key, coef_dict[key]) for key in matched_coef_cols])

    matched_coef_series = pd.Series(list(matched_coef_dict.values()), index=matched_coef_dict.keys())
    out_df.loc[:, matched_coef_cols] = out_df.loc[:, matched_coef_cols] * matched_coef_series
    return out_df,out_cond


if __name__ == '__main__':
    root_path=os.getcwd()
    finp = input('Enter file corresponding to "opt2D_R_Ref.txt ": \n').strip().split()
    lines_list=[]
    for finp_name  in finp:
        with open(os.path.join(root_path,finp_name)) as f:
            lines = f.readlines()
            lines_list.append(lines)
    if "Ref" in " ".join(finp):
        if "R_Ref" in " ".join(finp):
            print("可以匹配上 R_Ref")
            coef_config='Parse_Setfos_coef_R.json'
        if "G_Ref" in" ".join(finp):
            print("可以匹配上 G_Ref")
            coef_config = 'Parse_Setfos_coef_G.json'
        if "B_Ref" in " ".join(finp):
            print("可以匹配上 B_Ref")
            coef_config = 'Parse_Setfos_coef_B.json'
    else:
        coef_config = 'Parse_Setfos_coef.json'
    print(coef_config)
    out_df,cond_cols=main(lines_list,coef_config)
    fout = fout(finp[0], 'Setfos', "csv")
    out_df.to_csv(os.path.join(root_path, fout), index=False)
    print('%d Sweep Conditions are found in %s: %s' % (len(cond_cols), finp, ' '.join(cond_cols)))
    print('Total rows:' + TEXT_COLOR_GREEN + ' %d' % (len(out_df)) + TEXT_COLOR_END)
    print('Write output file ' + TEXT_COLOR_GREEN + "%s" % (fout,) + TEXT_COLOR_END + ' successfully!')

